#include "KD.h" 
#include <algorithm> 
#include <iostream>
#include <cassert>
#include <queue>
#include "sample1.h"
 
using namespace std;
using namespace optix; 


int max_component_idx(float3 v) {
	if (v.x > v.y) {
		if (v.x > v.z) {
			return 0;
		}
		else {
			return 2;
		}
	}
	else {
		if (v.y > v.z) {
			return 1;
		}
		else {
			return 2;
		}
	}
}

bool compare_x(const PhotonRec &lhs, const PhotonRec &rhs) {
	return lhs.pos.x < rhs.pos.x;
}
bool compare_y(const PhotonRec &lhs, const PhotonRec &rhs) {
	return lhs.pos.y < rhs.pos.y;
}
bool compare_z(const PhotonRec &lhs, const PhotonRec &rhs) {
	return lhs.pos.z < rhs.pos.z;
}

int KD::computeSize(uint pl_width, uint pl_height) {
	const int MPPR = Photon::maxPhotonsPerRay;
	int maxPhotons = pl_width*pl_height*MPPR;
	int pow=1;
	int size=0;
	for (; pow<maxPhotons; size+=pow, pow*=2);

	return size;
}

KD::KD (PhotonRec *photons, 
		uint pl_width, 
		uint pl_height, 
		float3 min_box, 
		float3 max_box) {  

	const int MPPR = Photon::maxPhotonsPerRay;
	maxPhotons = pl_width*pl_height*MPPR;
	size=computeSize(pl_width, pl_height);
	kd.resize(size);  

	vector<PhotonRec> vec;
	vec.reserve(maxPhotons);
	for (int i=0; i<pl_width; ++i) {
		for (int j=0; j<pl_height; ++j) {
			for (int k=0; k<MPPR; ++k) {
					PhotonRec &rec = photons[i*pl_height*MPPR + j*MPPR + k];
					if ( fminf(rec.pos) < fminf(min_box) || fmaxf(rec.pos) > fmaxf(max_box) ) {
						cout << "Photon out of bounds" << endl;
						exit(-1);
					}
					if (rec.energy.x > 0.f) {
						vec.push_back(rec);
					} 

			}
		}
	}

	build(vec, 0, 0, (int)vec.size()-1, 0, min_box, max_box);
}
 

void KD::build( vector<PhotonRec> &pts, 
				int ind,
				int start, 
				int last, 
				int depth,
				// To find the longest dimension
				float3 bbox_min, 
				float3 bbox_max) 
{ 
	if (ind >= size) {
		cout << "KD out of bounds" << endl;
		return;
	}

	if (start > last) {
		kd[ind].axis = -1; // flag as non-existent
		return; 
	}
	
	if (last == start) {
		kd[ind] = pts[start];
		kd[ind].axis = 3;
	}
	else {
		int split_dim = getSplitDim(bbox_min, bbox_max);
		int median_idx = getMedian(bbox_min, bbox_max, pts, start, last);
		kd[ind] = pts[median_idx];
		kd[ind].axis = split_dim;
		 
		// Build children 
		float3 new_bbox_max = bbox_max;
		float3 new_bbox_min = bbox_min;
		
		switch(split_dim) {
		case 0:
			new_bbox_max.x = new_bbox_min.x = kd[ind].pos.x;
			break;
		case 1:
			new_bbox_max.y = new_bbox_min.y = kd[ind].pos.y;
			break;
		case 2:
			new_bbox_max.z = new_bbox_min.z = kd[ind].pos.z;
			break;
		}
		  
		build( pts, ind*2 + 1, start, median_idx-1, depth+1, bbox_min, new_bbox_max );
		build( pts, ind*2 + 2, median_idx+1, last, depth+1, new_bbox_min, bbox_max );
	} 
} 
 
int KD::getMedian(float3 &bbmin, float3 &bbmax, vector<PhotonRec> &pts, int start, int last) {
	// Find longest dimension 
	int split_dim = getSplitDim(bbmin, bbmax);

	// Find median
	int median_idx = (start+last)/2;

	// Nth element 
	switch(split_dim) {
	case 0:
		nth_element(pts.begin() + start, 
					pts.begin() + median_idx,
					pts.begin() + last + 1,
					compare_x);
		break;
	case 1: 
		nth_element(pts.begin() + start, 
					pts.begin() + median_idx,
					pts.begin() + last + 1,
					compare_y);
		break;
	case 2: 
		nth_element(pts.begin() + start, 
					pts.begin() + median_idx,
					pts.begin() + last + 1,
					compare_z);
		break;
	default: ;
	}
	return median_idx; 
} 

int KD::getSplitDim(float3 &bbmin, float3 &bbmax) { 
	return max_component_idx(bbmax - bbmin);	
}
 
void KD::store( PhotonRec *buffer ) {
	for (int i=0; i<kd.size(); ++i) {
		buffer[i] = kd[i];
	}
}


void KD::gatherCPU( ViewRayData* hit_buffer, float* output, int width, int height, int plw, int plh ) {
	
	uint2 li;
	for (li.x=0; li.x < width; ++li.x) { 	 
		for (li.y=0; li.y < height; ++li.y) {

			ViewRayData hit = hit_buffer[li.x * height + li.y];
			PhotonRec photon;

			int mdep = 22;
			int stack[22];
			int top = 0, pid;
			stack[top] = 0;
			float r2 = 400.f;
			float d;
			float3 accum_flux = make_float3(0.0f);
			float3 diff;
			
			while ( top >= 0 ) {
				pid = stack[top--]; 
				photon = kd[pid];
				if ( photon.axis < 0 ) {
					continue;
				}
				diff = photon.pos - hit.position;
				d = length(diff);
				if (d*d < r2) {
					if (dot(photon.normal, hit.normal) > 0.001f) {
						accum_flux += photon.energy * hit.color;
					}
				}

				switch ( photon.axis ) {				
				case 0:
					d = diff.x;
					break;
				case 1:
					d = diff.y;
					break;
				case 2:
					d = diff.z;
					break;
				} 
				 
				if ( d*d < r2 ) { 
									
					if ( top < mdep && photon.axis != 3 ) {
						stack[++top] = pid * 2 + 1;
						stack[++top] = pid * 2 + 2;
					}
				}	

				else {
					int child = d < 0.0f ? 1 : 0;
					if ( top < mdep && photon.axis != 3 ) {
						stack[++top] = pid * 2 + child + 1;
					}
				}	 
			
			} //while - kd search
			 
 
			float3 indirect_light = 1.0f / ( M_PIf * r2 ) * accum_flux / (plw * plh / 2);
			output[li.x*height*4 + li.y*4 + 0] = indirect_light.x;
			output[li.x*height*4 + li.y*4 + 1] = indirect_light.y;
			output[li.x*height*4 + li.y*4 + 2] = indirect_light.z;
			output[li.x*height*4 + li.y*4 + 3] = 0.f;
		}

	}

}	